//
// Created by Never on 2019/10/12.
//

#ifndef MINIMAL_SPANNING_FOREST_H
#define MINIMAL_SPANNING_FOREST_H

const int unConnected = INT32_MIN;

template <typename val>
class MinimalSpanningForest{
protected:
    std::vector<bool> isKnown;
    std::vector<val> minCost;
    std::vector<int> preNode;
    std::vector<int> nodes;
    std::vector<std::vector<val>> distances;
    const int notFound = -1;
    val cost;

    bool isConnected(int node1, int node2);
    virtual void initialTable(std::vector<int> usedRoots);

    virtual void initialTableToDefault();
    int findNodeToInsert();
    virtual void updateMinCostAfterInsertingNode(int node);
    virtual bool isKnownNodeToNodeCloser(int knownNode, int node);
    void updateClosestKnownNodeOfNode(int node, int knownNode);
    virtual void handleNewInsertNode(int node);
    virtual bool isExit(int node);
public:
    MinimalSpanningForest(std::vector<int> roots, std::vector<int> nodes, std::vector<std::vector<val>> distances):
    nodes(nodes), distances(distances),
    isKnown(roots.size()+nodes.size(), false),
    minCost(roots.size()+nodes.size(), INT32_MAX),
    preNode(roots.size()+nodes.size(), 0), cost(0)
    {}
    val minimalSpanningForest(std::vector<int> usedRoots);
    val getCost();
};

template<typename val>
bool MinimalSpanningForest<val>::isConnected(int node1, int node2) {
    return distances[node1][node2] != unConnected;
}

template<typename val>
void MinimalSpanningForest<val>::initialTable(std::vector<int> usedRoots) {
    initialTableToDefault();
    for(auto root: usedRoots)
    {
        isKnown[root] = true;
        updateClosestKnownNodeOfNode(root, root);
        for(auto node: nodes)
        {
            if(isKnownNodeToNodeCloser(root, node))
                updateClosestKnownNodeOfNode(node, root);
        }
    }
}

template<typename val>
void MinimalSpanningForest<val>::initialTableToDefault() {
    for(int i = 0; i < isKnown.size(); i++)
    {
        isKnown[i] = false;
        minCost[i] = INT32_MAX;
        preNode[i] = 0;
    }
}

template<typename val>
int MinimalSpanningForest<val>::findNodeToInsert() {
    int nodeWithMinimalCost = notFound;
    val minimalCost = INT32_MAX;
    for(auto node: nodes)
    {
        if(not isKnown[node])
        {
            if(minCost[node] < minimalCost)
            {
                minimalCost = minCost[node];
                nodeWithMinimalCost = node;
            }
        }
    }
    return nodeWithMinimalCost;
}

template<typename val>
void MinimalSpanningForest<val>::updateMinCostAfterInsertingNode(int node) {
    for(auto n:nodes)
    {
        if(not isKnown[n])
        {
            if(isKnownNodeToNodeCloser(node, n))
                updateClosestKnownNodeOfNode(n, node);
        }
    }
}

template<typename val>
val MinimalSpanningForest<val>::minimalSpanningForest(std::vector<int> usedRoots) {
    cost = 0;
    initialTable(usedRoots);
    int node;
    while(isExit(node = findNodeToInsert()) )
        handleNewInsertNode(node);
    return cost;
}

template<typename val>
bool MinimalSpanningForest<val>::isKnownNodeToNodeCloser(int knownNode, int node) {
    return isConnected(knownNode, node) and minCost[node] > distances[knownNode][node];
}

template<typename val>
void MinimalSpanningForest<val>::updateClosestKnownNodeOfNode(int node, int knownNode) {
    minCost[node] = distances[knownNode][node];
    preNode[node] = knownNode;
}

template<typename val>
void MinimalSpanningForest<val>::handleNewInsertNode(int node) {
    cost += minCost[node];
    isKnown[node] = true;
    updateMinCostAfterInsertingNode(node);
}

template<typename val>
val MinimalSpanningForest<val>::getCost() {
    return cost;
}

template<typename val>
bool MinimalSpanningForest<val>::isExit(int node) {
    return node != notFound;
}

#endif //MINIMAL_SPANNING_FOREST_H
